"""辅助脚本：把 try.py 的主流程封装为函数，供 `main.py` 调用。

将实验性逻辑集中在这里，`try.py` 仅作为壳子调用本函数。
"""
import os
import time
import yaml
from matplotlib.font_manager import FontProperties

# 样式/字体工具
from utils.plot_style import set_plot_style, set_chinese_font

# core 与 plot 模块
from src.core_utils import (
	generate_apc_data,
	preprocess_data,
	print_descriptive_table,
	print_apc_identification_note,
	build_and_sample_apc_model,
	build_and_sample_full_model,
	build_and_sample_full_model_with_interaction,
	model_diagnostics,
	print_full_model_diagnosis,
	save_param_table,
	print_beta_age_edu,
	gender_group_analysis,
)
from src.plot import (
	explore_and_save_figures,
	generated_images,
	plot_and_save_trace,
	plot_effect,
	compare_effects,
)


def run_try_workflow():
	"""Run the experimental try workflow (moved here from try.py).

	This function replicates `try.py`'s `try_main` behavior so external callers
	(for example `main.py`) can invoke the workflow programmatically.
	"""
	print("=" * 40)
	start_time_str = time.strftime("%Y-%m-%d %H:%M:%S")
	start_time = time.time()
	print(f"⏱ APC开始 {start_time_str}")
	print("=" * 40)

	# 读取 config.yaml（若不存在则使用默认设置）
	cfg = {}
	try:
		with open("config.yaml", "r", encoding="utf-8") as f:
			cfg = yaml.safe_load(f) or {}
	except FileNotFoundError:
		print("未找到 config.yaml，使用内置默认配置。")
	except Exception as e:
		print("读取 config.yaml 时出错：", e)

	plot_cfg = cfg.get("plot") if cfg else None
	set_plot_style(plot_cfg)
	if plot_cfg and plot_cfg.get("font_family"):
		set_chinese_font(plot_cfg.get("font_family"))

	# 为绘图对象创建 FontProperties（部分函数需要 font 对象）
	font_family = (plot_cfg.get("font_family") if plot_cfg else None) or "Heiti TC"
	try:
		font_prop = FontProperties(family=font_family)
	except Exception:
		font_prop = None

	# 配置（与早前脚本保持兼容）
	config = cfg.get("simulation") or cfg.get("config") or cfg or {
		"n_ages": 10,
		"n_periods": 5,
		"n_cohorts": 10,
		"n_individuals_per_group": 200,
		# 下面采样相关参数默认关闭/较小，用于快速 smoke-test
		"draws": 100,
		"tune": 100,
		"cores": 1,
		"chains": 1,
		"random_seed": 42,
		"fig_save_path": cfg.get("figures_dir") or cfg.get("figure_path") or "figures",
		# 开关：是否执行耗时采样，默认 False（快速测试）
		"do_sampling": False,
	# 高级检查开关：模型比较、保存 InferenceData、分组 PPC、先验敏感性
	"run_advanced_checks": cfg.get("run_advanced_checks", False),
	}

	# 确保 figures 目录存在
	fig_dir = config.get("fig_save_path") or "figures"
	os.makedirs(fig_dir, exist_ok=True)

	# 1) 数据生成与预处理
	print("\n--- 数据生成与预处理 ---")
	data = generate_apc_data(config)
	data = preprocess_data(data)
	print("数据样例：")
	print(data.head())

	# 2) 描述性统计与说明
	print_descriptive_table(data)
	print_apc_identification_note()

	# 3) 探索性绘图（快速）
	print("\n--- 探索性绘图（保存到 figures）---")
	explore_and_save_figures(data, font_prop, fig_save_path=fig_dir, plot_cfg=plot_cfg)

	# 4) 建模与采样（不跳过）
	print("\n--- 建模与采样（可能较慢）---")
	trace_apc, n_ages, n_periods, n_cohorts, age_idx, period_idx, cohort_idx, y = build_and_sample_apc_model(data, config)
	trace_full = build_and_sample_full_model(data, n_ages, n_periods, n_cohorts, age_idx, period_idx, cohort_idx, y, config)
	trace_full_inter = build_and_sample_full_model_with_interaction(
		data, n_ages, n_periods, n_cohorts, age_idx, period_idx, cohort_idx, y, config
	)

	# 5) 模型诊断与检验
	model_diagnostics(trace_apc, "APC模型")
	model_diagnostics(trace_full, "完整模型")
	print_full_model_diagnosis(trace_full)

	# 6) 参数表格（终端输出）
	save_param_table(
		trace_full,
		['mu', 'age_effect', 'period_effect', 'cohort_effect', 'beta_gender', 'beta_region', 'beta_education']
	)

	# 7) 可视化与保存
	plot_and_save_trace(trace_apc, 'APC模型参数采样轨迹图', "apc_trace.png")
	plot_and_save_trace(trace_full, '完整模型参数采样轨迹图', "full_trace.png")
	plot_effect(trace_full, 'age_effect', n_ages, data['age_group'].cat.categories, '完整模型年龄效应估计 (95% HDI)', "full_age_effect.png", font_prop)
	import numpy as _np
	plot_effect(trace_full, 'period_effect', n_periods, _np.unique(data['period']), '完整模型时期效应估计 (95% HDI)', "full_period_effect.png", font_prop)
	plot_effect(trace_full, 'cohort_effect', n_cohorts, _np.unique(data['cohort_code']), '完整模型队列效应估计 (95% HDI)', "full_cohort_effect.png", font_prop)
	compare_effects(
		trace_apc, trace_full, 'age_effect', n_ages, data['age_group'].cat.categories,
		'年龄效应对比（APC模型 vs 完整模型）', "compare_age_effect.png", font_prop
	)

	# 8) 交互项与分组分析
	print_beta_age_edu(trace_full_inter)
	gender_group_analysis(data, config, font_prop)

	# 9) 高级检查（可选）
	if config.get("run_advanced_checks"):
		from src.core_utils import model_compare, save_inferencedata, group_ppc_plot, prior_sensitivity_run
		# 保存 InferenceData
		save_inferencedata(trace_apc, "trace_apc", out_dir=config.get("output_path", "outputs"))
		save_inferencedata(trace_full, "trace_full", out_dir=config.get("output_path", "outputs"))
		save_inferencedata(trace_full_inter, "trace_full_inter", out_dir=config.get("output_path", "outputs"))

		# 模型比较
		cmp = model_compare({"APC": trace_apc, "Full": trace_full, "FullInter": trace_full_inter})

		# 分组 PPC（按 gender）
		try:
			group_ppc_plot(trace_full, data, "gender", var_name="y_obs", fig_save_path=fig_dir, plot_cfg=plot_cfg)
		except Exception as e:
			print("分组 PPC 失败：", e)

		# 先验敏感性（示例，需模型 builder 支持 cfg['prior_std']）
		try:
			def builder(d, c):
				# 简单示例：调用 full model builder，期望其读取 c['prior_std']（当前未使用）
				return build_and_sample_full_model(d, n_ages, n_periods, n_cohorts, age_idx, period_idx, cohort_idx, y, c)

			sens = prior_sensitivity_run(data, config, prior_stds=[10, 25, 50], model_builder=builder, var_of_interest="beta_education")
			print("先验敏感性结果：", sens)
		except Exception as e:
			print("先验敏感性分析失败：", e)

	# 展示并列出生成的图片（不弹窗）
	print("\n已生成图片：")
	for p in generated_images:
		print(p)

	end_time_str = time.strftime("%Y-%m-%d %H:%M:%S")
	end_time = time.time()
	print("=" * 40)
	print(f"⌛️ APC结束 {end_time_str}")
	duration = end_time - start_time
	print(f"总耗时：{int(duration)} 秒（{duration/60:.2f} 分钟）")
	print("=" * 40)

	return {
		'status': 'done',
		'start': start_time_str,
		'end': end_time_str,
	}

